"""
处理召回需要的数据
"""
import json

import pandas

import chatbot.lib as lib


def prepar_recall_corpus():
    q_path = "corpus/dnn/recall/A.txt"
    a_path = "corpus/dnn/recall/Q.txt"
    qa_path = "corpus/dnn/recall/QA.txt"
    excel_path = "corpus/dnn/recall/excel.xlsx"

    q_lines = open(q_path, encoding="utf-8").readlines()
    a_lines = open(a_path, encoding="utf-8").readlines()
    qa_dict = {}

    for q, a in zip(q_lines, a_lines):
        qa_dict[q.strip()] = {}
        qa_dict[q.strip()]["answer"] = a.strip()
        qa_dict[q.strip()]["q_cut_by_word"] = lib.cut(q, by_word=True)
        result = lib.cut(q, with_sg=True)
        qa_dict[q.strip()]["q_cut"] = [i[0] for i in result]
        entity = [i[0] for i in result if i[1] == "kc"]
        qa_dict[q.strip()]["entity"] = entity

    df = pandas.read_excel(excel_path)[["问题", "答案"]]
    for q, a in zip(df["问题"], df["答案"]):
        qa_dict[q.strip()] = {}
        qa_dict[q.strip()]["answer"] = a.strip()
        qa_dict[q.strip()]["q_cut_by_word"] = lib.cut(q, by_word=True)
        result = lib.cut(q, with_sg=True)
        qa_dict[q.strip()]["q_cut"] = [i[0] for i in result]
        entity = [i[0] for i in result if i[1] == "kc"]
        qa_dict[q.strip()]["entity"] = entity

    json.dump(qa_dict, open(qa_path, "w"), ensure_ascii=False, indent=2)
